﻿using System.Numerics;

namespace LiteVectorIndex.Engine
{
    public static class DistanceProvider
    {
        public static float Similarity(ReadOnlySpan<float> v1, ReadOnlySpan<float> v2, MetricType metric)
        {
            switch (metric)
            {
                case MetricType.L2:
                    return 1.0f / (1.0f + L2DistanceSimd(v1, v2)); // L2距离的相似性
                case MetricType.InnerProduct:
                    return DotSimd(v1, v2); // 内积
                case MetricType.Cosine:
                    // 使用临时数组做归一化，不修改原始数据
                    float[] v1Norm = new float[v1.Length];
                    float[] v2Norm = new float[v2.Length];
                    NormalizeInto(v1, v1Norm);
                    NormalizeInto(v2, v2Norm);
                    return DotSimd(v1Norm, v2Norm); // 归一化后的点积即余弦
                default:
                    throw new NotImplementedException();
            }
        }

        // 原始归一化写到新数组，不修改输入
        public static void NormalizeInto(ReadOnlySpan<float> src, Span<float> dst)
        {
            float norm = Norm(src);
            if (norm < 1e-10f)
            {
                for (int i = 0; i < dst.Length; i++) dst[i] = 0f;
                return;
            }
            for (int i = 0; i < src.Length; i++)
                dst[i] = src[i] / norm;
        }

        private static float L2DistanceSimd(ReadOnlySpan<float> v1, ReadOnlySpan<float> v2)
        {
            float sum = 0f;
            int i = 0;
            int simdLength = Vector<float>.Count;

            for (; i <= v1.Length - simdLength; i += simdLength)
            {
                var a = new Vector<float>(v1.Slice(i));
                var b = new Vector<float>(v2.Slice(i));
                var diff = a - b;
                sum += Vector.Dot(diff, diff);
            }

            for (; i < v1.Length; i++)
                sum += (v1[i] - v2[i]) * (v1[i] - v2[i]);

            return MathF.Sqrt(sum);
        }

        private static float DotSimd(ReadOnlySpan<float> v1, ReadOnlySpan<float> v2)
        {
            float dot = 0f;
            int i = 0;
            int simdLength = Vector<float>.Count;

            for (; i <= v1.Length - simdLength; i += simdLength)
            {
                var a = new Vector<float>(v1.Slice(i));
                var b = new Vector<float>(v2.Slice(i));
                dot += Vector.Dot(a, b);
            }

            for (; i < v1.Length; i++)
                dot += v1[i] * v2[i];

            return dot;
        }

        private static float Norm(ReadOnlySpan<float> v)
        {
            float sum = 0f;
            foreach (var x in v)
                sum += x * x;
            return MathF.Sqrt(sum);
        }
    }
}